#include "nets.hpp"
#include "ATen/ops/relu.h"
#include "torch/nn/modules/linear.h"

Net::Net()
{
    const int num = 16;

    input_  = register_module("input", torch::nn::Linear(1, num));
    hidden_ = register_module("hidden", torch::nn::Linear(num, num));
    output_ = register_module("output", torch::nn::Linear(num, 1));
}

Net::~Net() {}

at::Tensor Net::forward(at::Tensor x)
{
    x = input_->forward(x);
    x = torch::relu(x);

    x = hidden_->forward(x);
    x = torch::relu(x);

    x = output_->forward(x);
    return x;
}

CustomDtaset::CustomDtaset(at::Tensor input_data, at::Tensor label_data)
{
    inputs_ = input_data;
    labels_ = label_data;
}

torch::data::Example<> CustomDtaset::get(size_t index)
{
    torch::Tensor Input = inputs_[index];
    torch::Tensor Label = labels_[index];
    return {Input.clone(), Label.clone()};
}

torch::optional<size_t> CustomDtaset::size() const { return labels_.size(0); }
